from PIL import Image
from skimage.morphology import erosion, dilation
import os
import numpy as np
from tqdm import tqdm
import cv2

def gaussian(image):
    kernel_size = 3
    sigmaX = 0
    # 如果不指定 sigmaX，则由 kernel_size 自动计算
    blurred_image = cv2.GaussianBlur(image, (kernel_size, kernel_size), sigmaX)#高斯滤波
    median_filtered_image = cv2.medianBlur(image, kernel_size)#中值滤波
    return median_filtered_image

def black_process(image):
    # 转换为灰度图像
    gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    # 应用阈值操作
    _, binary_image = cv2.threshold(gray_image, 230, 255, cv2.THRESH_BINARY)
    # 寻找水印区域（这里需要根据图像调整条件）
    white_areas = binary_image > 200  # 假设水印是较亮的白色
    white_areas = white_areas.astype(np.uint8)
    # 使用膨胀操作来连接水印区域内的点
    kernel = np.ones((5, 5), np.uint8)
    watermark_mask = cv2.dilate(white_areas, kernel, iterations=2)
    # 用周围像素的平均值填充水印区域
    inverted_mask = cv2.bitwise_not(watermark_mask)
    original_image_without_watermark = cv2.bitwise_and(image, image, mask=inverted_mask)
    filled_image = cv2.inpaint(original_image_without_watermark, watermark_mask, 3, cv2.INPAINT_TELEA)
    return filled_image

def white_process(image):
    # 转换为灰度图像
    gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    # 应用阈值操作
    _, binary_image = cv2.threshold(gray_image, 0, 80, cv2.THRESH_BINARY)
    # 寻找水印区域（这里需要根据图像调整条件）
    white_areas = binary_image < 80  # 假设水印是较亮的白色
    white_areas = white_areas.astype(np.uint8)
    # 使用膨胀操作来连接水印区域内的点
    kernel = np.ones((5, 5), np.uint8)
    watermark_mask = cv2.dilate(white_areas, kernel, iterations=2)
    # 用周围像素的平均值填充水印区域
    inverted_mask = cv2.bitwise_not(watermark_mask)
    original_image_without_watermark = cv2.bitwise_and(image, image, mask=inverted_mask)
    filled_image = cv2.inpaint(original_image_without_watermark, watermark_mask, 3, cv2.INPAINT_TELEA)
    return filled_image

# 设置图像素材文件夹和形态学处理结果文件夹
input_folder = '新图像素材/test'
output_folder = '新预处理/test'

# 确保输出文件夹存在
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 定义开运算的结构元素，这里使用方形结构元素
# 参数为结构元素的大小，例如(3, 3)表示3x3的结构元素


# 对文件夹中的所有图片进行处理
for n, filename in enumerate(sorted(os.listdir(input_folder)), start=1):
    if filename.lower().endswith(('.png', '.jpg', '.bmp', '.gif')):
        # 构造完整的文件路径
        image_path = os.path.join(input_folder, filename)
        # 打开图像
        image = Image.open(image_path)  # 转换为灰度图
        # 将PIL图像转换为numpy数组
        image_array = np.array(image)
        print(image_array.shape)

        # 计算平均亮度
        average_luminance = np.mean(image_array)
        # 判断背景颜色
        if average_luminance > 128:  # 阈值可以根据需要调整
            background_color = 'white'
        else:
            background_color = 'black'

        h = image_array.shape[0]
        w = image_array.shape[1]

        #剪裁
        if h>w:
            width, height = image.size
            left = int(width * 0.2)
            upper = int(height * 0.11)
            right = int(width * 0.8)  # 宽度减去左边裁剪的部分
            lower = int(height * 0.5) # 初始高度，因为我们还没有裁剪底部
            cropped_image = image.crop((left, upper, right, lower))
            cropped_image_array = np.array(cropped_image)
            hc = cropped_image_array.shape[0]
            wc = cropped_image_array.shape[1]
            a = int(8.8*(hc//10))
            if background_color == 'white':
                for i in range(a,hc):
                    for j in range(wc):
                        if cropped_image_array[i,j,0] < 150 and cropped_image_array[i,j,1] < 150 and cropped_image_array[i,j,2] < 150:
                            cropped_image_array[i, j, 0] = 255
                            cropped_image_array[i, j, 1] = 255
                            cropped_image_array[i, j, 2] = 255
            else:
                for i in range(9*(hc//10), hc):
                    for j in range(wc):
                        if cropped_image_array[i,j,0] > 150 and cropped_image_array[i,j,1] > 150 and cropped_image_array[i,j,2] > 150:
                            cropped_image_array[i, j, :] = 0
        else:
            if background_color == 'black':
                cropped_image_array = black_process(image_array)
            else:
                cropped_image_array = image_array
        if cropped_image_array.shape[2]==4:
            cropped_image_array = cropped_image_array[:,:,:3]
        cropped_image = Image.fromarray(cropped_image_array)
        '''if background_color == 'black':
            cropped_image_array = black_process(image_array)
            #cropped_image_array = gaussian(cropped_image_array)
        else:
            #cropped_image_array = white_process(image_array)
            #cropped_image_array = gaussian(cropped_image_array)
            cropped_image_array = image_array
        cropped_image = Image.fromarray(cropped_image_array)'''

        # 构造新的文件名，按照12345的顺序
        base_filename = os.path.splitext(filename)[0]
        new_filename = f'{base_filename}.jpg'
        new_filepath = os.path.join(output_folder, new_filename)

        # 保存处理后的图像
        cropped_image.save(new_filepath)
        print(f'处理后的图像已保存：{new_filepath}')


